import os

import pandas as pd

from tools.df_data_utils import is_none_null


def get_all_dict_from_df_data(df_data, key_column_list, value_column_list):
    is_list = True
    # 如果key_column_list或value_column_list不是列表，则转换为列表
    if not isinstance(key_column_list, list):
        is_list = False
        key_column_list = [key_column_list]
    if not isinstance(value_column_list, list):
        value_column_list = [value_column_list]

    # 确保键值对列的数量匹配
    if len(key_column_list) != len(value_column_list):
        raise ValueError("Key columns and value columns must have the same number of elements.")

    # 初始化字典列表，每个键值对列组合一个字典
    data_dict_list = [{} for _ in key_column_list]

    def my_row_func(index, row):
        for i, (key_column, value_column) in enumerate(zip(key_column_list, value_column_list)):
            if key_column not in df_data.columns or value_column not in df_data.columns:
                continue  # 跳过不存在的列
            key = row[key_column]
            value = row[value_column]
            # 过滤掉键或值为空、空字符串或包含无效字符的情况
            if is_none_null(key) or is_none_null(value):
                continue
            # 向对应的字典中添加键值对
            data_dict_list[i][key] = value

    # 遍历DataFrame中的每一行
    for index, row in df_data.iterrows():
        my_row_func(index, row)

    if is_list:
        return data_dict_list
    else:
        return data_dict_list[0]


def get_all_dict_from_csv(csv_file_name, key_column_list, value_column_list):
    # 如果key_column_list或value_column_list为空，则返回空字典
    if not os.path.exists(csv_file_name):
        if isinstance(key_column_list, list):
            return [{} for _ in key_column_list]
        else:
            return {}
    all_dict_list = [{} for _ in key_column_list]
    chunksize = 10000
    # 读取文件数据
    for df_chunk in pd.read_csv(csv_file_name, chunksize=chunksize, encoding="utf-8", dtype=str):
        df_chunk_copy = df_chunk.copy()
        re_dict_list = get_all_dict_from_df_data(df_chunk_copy, key_column_list, value_column_list)
        for all_dict, re_dict in zip(all_dict_list, re_dict_list):
            all_dict.update(re_dict)
    return all_dict_list


def fill_missing_data_dict(df_data, my_dict, key_name_list,
                           value_name):
    if not isinstance(key_name_list, list):
        key_name_list = [key_name_list]
    if value_name not in df_data.columns:
        df_data[value_name] = ''

    def my_func(row):
        if not is_none_null(row[value_name]):
            return row[value_name]
        for key in key_name_list:
            if key in row:
                value = my_dict.get(row[key], '')
                if is_none_null(value):
                    continue
                return value
        return ''

    df_data[value_name] = df_data.apply(my_func, axis=1)


def fill_missing_data_set(df_data, my_set, key_name_list,
                          value_name):
    if not isinstance(key_name_list, list):
        key_name_list = [key_name_list]
    if value_name not in df_data.columns:
        df_data[value_name] = '否'

    def my_func(row):
        for key in key_name_list:
            if key in row:
                if row[key] in my_set:
                    return "是"
        return row[value_name]

    df_data[value_name] = df_data.apply(my_func, axis=1)
